#include <arch/cpu.h>
#include <mm/mm.h>
#include <mm/rmap.h>
#include <mm/valloc.h>
#include <mm/slab.h>
#include <mm/swap.h>
#include <ds/lru.h>
#include <ds/bitmap.h>
#include <sys/unistd.h>
#include <sys/sysdef.h>
#include <device/block.h>

#define LOG_TAG "kswapd"
#include <log.h>

int evict_page(lnode_t* node);

lru_t swap_lru = {
    .head = {
        .next = &swap_lru.head,
        .prev = &swap_lru.head
    },
    .size = 0,
    .lock = {
        .owner = NULL,
        .sem.count = 1
    },
    .destory = evict_page
};

extern uptr max_pg;

kmem_cache_t* swap_cache;

lnode_t swap_devices = { .next = &swap_devices, .prev = &swap_devices };
volatile u32 swap_id = 0;

void set_slot_free(struct swap_device* swap, pg_entry_t* pte) {
    if (!swap) {
        list_for_each_entry(swap, &swap_devices, list) {
            if (swap->id == pte->swap.swap_id) {
                goto found;
            }
        }
        return;
    }

found:
    swap->free_slots++;
    bitmap_clear(swap->swap_map, pte->swap.slot);
    if (pte->swap.slot < swap->lowest_bit) {
        swap->lowest_bit = pte->swap.slot;
    } else if (pte->swap.slot > swap->highest_bit) {
        swap->highest_bit = pte->swap.slot;
    }
    return;
}

int do_swap_in(pg_entry_t* pte, uptr vaddr) {
    struct swap_device* swap;
    list_for_each_entry(swap, &swap_devices, list) {
        if (swap->id == pte->swap.swap_id) {
            goto found;
        }
    }

    return EFAULT;

found:
    page_t* page = NULL;
    if (!bitmap_get(swap->swap_map, pte->swap.slot)) {
        struct swap_cache_info* cache = NULL;
        list_for_each_entry(cache, &swap->queue, list) {
            if (cache->info->slot == pte->swap.slot) {
                page = cache->page;
                cache->page = NULL;
                break;
            }
        }
        assert_msg(page != NULL, "swap cache not found");
    } else {
        uptr new_page = get_page(PG_LOCKED);
        if (!new_page) {
            return ENOMEM;
        }

        mapping_p2v(PG_MP_4, new_page, PG_WRITE, VMAP_NULL);
        
        swap->dev->read(swap->dev, (void*)PG_MP_4, pte->swap.slot * PG_SIZE, PG_SIZE);
        
        unmapping_p2v(PG_MP_4);

        page = page_query(new_page);
        set_slot_free(swap, pte);
    }

    page->index = vaddr >> PG_SIZE_BITS;
    page->ref_counts = 0;

    try_to_remap_anon(page);
    set_page_unlock(page);
    return 0;
}

int do_swap_out(page_t* page) {
    int errno;
    uptr pa = page_query_pn(page) << PG_SIZE_BITS;
    struct swap_device* swap;
    list_for_each_entry(swap, &swap_devices, list) {
        if (swap->free_slots != 0) {
            break;
        }
    }

    if (swap->free_slots == 0) {
        return 0;
    }
    // TODO: add file mapping support, A different execute path is needed
    u32 slot = bitmap_find_first_zero_from(swap->swap_map, swap->highest_bit, swap->lowest_bit);
    struct swap_info* info = valloc(sizeof(struct swap_info));
    *info = (struct swap_info){
        .swap_id = swap->id,
        .slot = slot
    };
    struct swap_cache_info* out_page = take_obj(swap_cache);
    out_page->page = page;
    out_page->info = info;
    list_append(&swap->queue, &out_page->list);

    if ((errno = try_to_unmap_anon(page, info))) {
        goto out;
    }

    mapping_p2v(PG_MP_4, pa, PG_WRITE, VMAP_NULL);

    swap->dev->write(swap->dev, (void*)PG_MP_4, slot * PG_SIZE, PG_SIZE);

    unmapping_p2v(PG_MP_4);
    cli();
    if (!out_page->page) {
        errno = 1;
        goto out;
    }
    bitmap_set(swap->swap_map, slot);
    swap->free_slots--;
    swap->lowest_bit = slot;
    if (slot > swap->highest_bit) {
        swap->highest_bit = slot;
    }

    page->ref_counts--;
    assert(page->ref_counts == 0);
    page->flags = 0;
    page->index = 0;
    page->mapping = NULL;

out:
    list_remove(&out_page->list);
    release_obj(swap_cache, out_page);

    vfree(info);
    sti();
    return errno;
}

int evict_page(lnode_t* node) {
    page_t* page = container_of(node, page_t, lru);

    if (page->ref_counts == 0 || page->mapping == NULL) {
        return 1;
    }

    int errno = do_swap_out(page);
    if (errno < 0) {
        LOGE("(%d)Failed to swap out page %p\n", errno, page->index << PG_SIZE_BITS);
    }

    return errno == 0;
}

void kswapd(void* data) {
    bool out = 0;
    u32 allocable_pages = 0;
    u32 start_watermark = max_pg >> 1;
    u32 aggressive_watermark = max_pg >> 2;
    page_t* page;
    while (1) {
        allocable_pages = 0;
        if (list_empty(&swap_devices) || swap_lru.size == 0) {
            goto next;
        }

        for (u32 i = 0; i < max_pg << PG_SIZE_BITS; i += PG_SIZE) {
            page = page_query(i);
            if (page->ref_counts == 0) {
                allocable_pages++;
            }
        }
        if (allocable_pages < aggressive_watermark) {
            lru_evict_count(&swap_lru, 15);
            // lru_evict_half(&swap_lru);
        } else if (allocable_pages < start_watermark) {
            lru_evict_one(&swap_lru);
        }

next:
        yield();
        if (out) {
            break;
        }
    }
    exit();
}

__DEF_SYSCALL1(int, swapon, const char*, path) {
    dentry_t* dentry;
    int errno;
    if ((errno = fs_path_walk(path, NULL, &dentry, NULL, 0))) {
        return errno;
    }
    device_t* dev = (device_t*)dentry->inode->data;
    blk_dev_t* bdev = dev->impl;
    if (bdev->part_id != 0x82) {
        return ENODEV;
    }

    if (!swap_cache) {
        swap_cache = new_kmem_cache("swap_cache", sizeof(struct swap_cache_info), 1, 0);
    }
    uptr* tmp = take_obj(swap_cache);
    release_obj(swap_cache, tmp);

    struct swap_device* info = vzalloc(sizeof(struct swap_device));
    info->dev = dev;
    info->pages = (bdev->end_lba - bdev->start_lba) * bdev->blk_size / PG_SIZE;
    info->free_slots = info->pages;
    info->lowest_bit = 0;
    info->highest_bit = info->pages - 1;
    info->swap_map = vzalloc(info->pages / 8);
    info->id = swap_id++;
    list_head_init(&info->queue);
    
    list_append(&swap_devices, &info->list);
    return 0;
}